#include <math.h>
#include <stdio.h>
#include <stdlib.h>

// Read a vector from a binary file.
// n
// b1
// b2
// ...
// bn
int VectorRead_Binary(const char* file_name, double** b, int* n)
{
    FILE* fp = NULL;

    double* vec = NULL;
    int     size;
    int     j;

    /*----------------------------------------------------------
     * Read in the data
     *----------------------------------------------------------*/

    fp = fopen(file_name, "rb");

    fread(&size, sizeof(int), 1, fp);

    if (size <= 0) {
        fprintf(stderr, "Error: File %s is empty or has invalid size.\n", file_name);
        fclose(fp);
        return -1;
    }

    vec = (double*)malloc(size * sizeof(double));

    fread(vec, sizeof(double), size, fp);

    // for (int i = 0; i < size; i++) {
    //     printf("vec[%d] =  %e\n", i, vec[i]);
    // }

    fclose(fp);

    // output
    *b = vec;
    *n = size;

    return 0;
}

/*--------------------------------------------------------------------------
 * Read a CSR matrix from a binary file.
 *   - row: number of rows
 *   - ia(i), i=0:row
 *   - ja(i), i=0:nnz-1
 *   - aa(i), i=0:nnz-1
 *--------------------------------------------------------------------------*/

int CSRMatrixRead_Binary(const char* file_name, int* row, int* col, int* nnz, int** ia, int** ja, double** aa)
{
    FILE* fp = NULL;

    double* matrix_data;
    int*    matrix_i;
    int*    matrix_j;
    int     num_rows;
    int     num_nonzeros;
    int     file_base = 0; // default is 0-based index, it automatically judged based on the file_name

    int j;

    /*----------------------------------------------------------
     * Read in the data
     *----------------------------------------------------------*/
    fp = fopen(file_name, "rb");

    fread(&num_rows, sizeof(int), 1, fp);

    // read IA
    matrix_i = (int*)malloc((num_rows + 1) * sizeof(int));

    fread(matrix_i, sizeof(int), num_rows + 1, fp);
    if (matrix_i[0] == 1) {
        file_base = 1;
        for (j = 0; j < num_rows + 1; j++) {
            matrix_i[j] -= file_base;
        }
    }

    num_nonzeros = matrix_i[num_rows];

    matrix_j    = (int*)malloc(num_nonzeros * sizeof(int));
    matrix_data = (double*)malloc(num_nonzeros * sizeof(double));

    // read JA
    fread(matrix_j, sizeof(int), num_nonzeros, fp);
    if (file_base == 1) {
        for (j = 0; j < num_nonzeros; j++) {
            matrix_j[j] -= file_base;
        }
    }

    // read val
    fread(matrix_data, sizeof(double), num_nonzeros, fp);

    fclose(fp);

    // Output
    *row = num_rows;
    *col = num_rows;
    *nnz = num_nonzeros;
    *ia  = matrix_i;
    *ja  = matrix_j;
    *aa  = matrix_data;

    return 0;
}

/*--------------------------------------------------------------------------
 * Read a COO matrix from a binary file.
 *   - global_num_rows: number of rows for the global matrix
 *   - global_num_cols: number of columns for the global matrix
 *   - num_rows: number of rows of the local matrix
 *   - num_cols: number of columns of the local matrix
 *   - num_nonzeros: number of nonzeros of the local matrix
 *   - I(k), J(k), val(k), k = 1,...,num_nonzeros, I and J are global indices, 1-based
 *--------------------------------------------------------------------------*/

int COOMatrixRead_Binary(const char* file_name, int* glo_row, int* glo_col, int* row, int* col, int* nnz, int** ii, int** jj, double** aa)
{
    FILE* fp = NULL;

    int     global_num_rows, global_num_cols;
    double* val;
    int*    I;
    int*    J;
    int     num_rows, num_cols;
    int     num_nonzeros;
    int     file_base = 1; // default is 1-based index

    int k;

    /*----------------------------------------------------------
     * Read in the data
     *----------------------------------------------------------*/
    fp = fopen(file_name, "rb");

    fread(&global_num_rows, sizeof(int), 1, fp);
    fread(&global_num_cols, sizeof(int), 1, fp);
    fread(&num_rows, sizeof(int), 1, fp);
    fread(&num_cols, sizeof(int), 1, fp);
    fread(&num_nonzeros, sizeof(int), 1, fp);

    // malloc memory for I, J and val;
    I   = (int*)malloc(num_nonzeros * sizeof(int));
    J   = (int*)malloc(num_nonzeros * sizeof(int));
    val = (double*)malloc(num_nonzeros * sizeof(double));

    // read IA
    for (k = 0; k < num_nonzeros; k++) {
        fread(&I[k], sizeof(int), 1, fp);
        fread(&J[k], sizeof(int), 1, fp);
        fread(&val[k], sizeof(double), 1, fp);
        I[k]--;
        J[k]--; // 0-based index
    }

    fclose(fp);

    // Output
    *glo_row = global_num_rows;
    *glo_col = global_num_cols;
    *row     = num_rows;
    *col     = num_cols;
    *nnz     = num_nonzeros;
    *ii      = I;
    *jj      = J;
    *aa      = val;

    return 0;
}

// store x to a file
void store_x(int n, double* x, char* filename)
{
    FILE* p = fopen(filename, "w");
    fprintf(p, "%d\n", n);
    for (int i = 0; i < n; i++) fprintf(p, "%lf\n", x[i]);
    fclose(p);
}

int CSRMatrixPrint(char* file_name, int row, int col, int nnz, int* ia, int* ja, double* aa)
{
    FILE* fp = NULL;

    int j;
    int file_base = 1;
    int ierr      = 0;

    /*----------------------------------------------------------
     * Print the matrix data
     *----------------------------------------------------------*/

    fp = fopen(file_name, "w");

    fprintf(fp, "%d\n", row);

    for (j = 0; j <= row; j++) {
        fprintf(fp, "%d\n", ia[j] + file_base);
    }

    for (j = 0; j < ia[row]; j++) {
        fprintf(fp, "%d\n", ja[j] + file_base);
    }

    for (j = 0; j < ia[row]; j++) {
        fprintf(fp, "%.15e\n", aa[j]);
    }

    fclose(fp);

    return ierr;
}
